#ifndef ARCH_X86_MEMORYMANAGEMENT_H
#define ARCH_X86_MEMORYMANAGEMENT_H
# ifdef LANG_X86

# include <Globals.h>
# include <MemoryManagement.h>
# include <MonitorDriver.h>
# include <Asm.h>

# define PAGE_SIZE 0x1000
# define PAGE_MASK 0xFFFFF000

/**
  A page is a pointer to a frame.
**/
class Page
{
public:
  Page() :
    present(0), rw(0), user(0), accessed(0), dirty(0), 
    unused(0), frame(0) {}
  
  /**
    Standard get/set functions
  **/
  bool getPresent()  {return present == 1;}
  bool getRw()       {return rw == 1;}
  bool getUser()     {return user == 1;}
  bool getAccessed() {return accessed == 1;}
  bool getDirty()    {return dirty == 1;}
  u32int getFrame()  {return frame<<12;}
  void setPresent(bool b) { (b)?present=1:present=0; }
  void setRw(bool b) { (b)?rw=1:rw=0; }
  void setUser(bool b) { (b)?user=1:user=0; }
  void setAccessed(bool b) { (b)?accessed=1:accessed=0; }
  void setDirty(bool b) { (b)?dirty=1:dirty=0; }
  void setFrame(u32int f) { frame=(f>>12) /*& 0xFFFFF000*/; }
private:
  u32int present    : 1;   // Page present in memory
  u32int rw         : 1;   // Read-only if clear, readwrite if set
  u32int user       : 1;   // Supervisor level only if clear
  u32int accessed   : 1;   // Has the page been accessed since last refresh?
  u32int dirty      : 1;   // Has the page been written to since last refresh?
  u32int unused     : 7;   // Amalgamation of unused and reserved bits
  u32int frame     : 20;  // Frame address (shifted right 12 bits)
};

/**
  A page table holds 1024 pages
**/
class PageTable
{
public:
  PageTable()
  {
    for(int i=0;i<1024;i++)
    {
      pages[i] = Page();
    }
  }

  Page *getPage(u32int n)
  {
    return &pages[n];
  }

  PageTable *clone(Address *phys)
  {
    PageTable *table = new(true, phys) PageTable();
    for(int i = 0; i < 1024; i++)
    {
      if (pages[i].getFrame())
      {
        memoryManager.allocFrame(&table->pages[i]);
        if (pages[i].getPresent()) {table->pages[i].setPresent(true);}
        if (pages[i].getRw()) {table->pages[i].setRw(true);}
        if (pages[i].getUser()) {table->pages[i].setUser(true);}
        if (pages[i].getAccessed()) {table->pages[i].setAccessed(true);}
        if (pages[i].getDirty()) {table->pages[i].setDirty(true);}
        copyPagePhysical(pages[i].getFrame(), table->pages[i].getFrame());
      }
    }
    return table;
  }

private:
  Page pages[1024];
};

/**
  Holds 1024 pagetables.
**/
class PageDirectory
{
public:
  PageDirectory()
  {
    for (int i = 0; i < 1024; i++)
    {
      tables[i] = 0;
      tablesPhysical[i] = 0;
    }
    physicalAddr = (Address)tablesPhysical;
  }

  PageTable *getTable(u32int index)
  {
    return tables[index];
  }

  Address getPhysical()
  {
    return physicalAddr;
  }

  Page *getPage(Address addr, bool assign=true)
  {
    addr /= 4096;
    u32int tableIdx = addr/1024;
    if (tables[tableIdx]) // if the table is already assigned
    {
      return tables[tableIdx]->getPage(addr%1024);
    }
    else if(assign)
    {
      u32int tmp;
      tables[tableIdx] = new(true/*page aligned*/, &tmp/*give phys. addr */) PageTable();
      tablesPhysical[tableIdx] = tmp | 0x7; // PRESENT, RW, US
      return tables[tableIdx]->getPage(addr%1024);
    }
    else
    {
      return NULL;
    }
  }
  
  /**
    Create a new page directory, that is an identical copy to 'dir'.
    TODO: implement copy-on-write.
  **/
  PageDirectory *clone()
  {
    Address phys;
    PageDirectory *dir = new(true, &phys) PageDirectory();

    // Get the offset of tablesPhysical from the start of 
    // the PageDirectory object.
    Address offset = (Address)dir->tablesPhysical - (Address)dir;
    
    // Then the physical address of dir->tablesPhysical is
    dir->physicalAddr = phys + offset;

    // go through each page table. If the page table is in the
    // kernel directory, do not make a new copy.
    for (int i = 0; i < 1024; i++)
    {
      if (!tables[i]) continue;
      if (memoryManager.getKernelDirectory()->getTable(i) == tables[i])
      {
        // It's in the kernel, so just use the same pointer.
        dir->tables[i] = tables[i];
        dir->tablesPhysical[i] = tablesPhysical[i];
      }
      else
      {
        // copy the table.
        Address phys;
        dir->tables[i] = tables[i]->clone(&phys);
        dir->tablesPhysical[i] = phys | 0x07;
      }
    }
    
    return dir;
  }
  
  void dump()
  {
    bool b = false;
    kerr.write((char*)"Dumping page directory\n");
    for (int i = 0; i < 0xFFFFF; i++)
    {
      Page *page;
      if (tables[i/1024])
        page = tables[i/1024]->getPage(i%1024);
      else if (b)
      {
        kerr.write("\tEnd of table.\n");
        b = false;
      }
      else
        continue;
        
      if (!b && page->getPresent())
      {
        b = true;
        kerr.write("\t");
        kerr.writeHex(i<<12);
        kerr.write("\n\t\t.\n");
      }
      else if(b && !page->getPresent())
      {
        b = false;
        kerr.write("\t");
        kerr.writeHex(i<<12);
        kerr.write("\n");
      }
    }
  }

//private:
  /**
    Array of pointers to pagetables.
  **/
  PageTable *tables[1024];
  
  /**
    Array of pointers to the pagetables above, but gives their *physical*
    location, for loading into the CR3 register.
  **/
  u32int tablesPhysical[1024];
  
  /**
    The physical address of tablesPhysical.
  **/
  Address physicalAddr;
};

# endif // LANG_X86
#endif  // ARCH_X86_MEMORYMANAGEMENT_H
